package pass

import (
	"crypto/rand"
	"crypto/subtle"
	"encoding/base64"
	"fmt"
	"golang.org/x/crypto/argon2"
	"io"
	"strings"
)

type argon2Params struct {
	memory      uint32
	iterations  uint32
	parallelism uint8
	saltLength  uint32
	keyLength   uint32
}

func compare(a, b *argon2Params) bool {
	return a.memory == b.memory &&
		a.iterations == b.iterations &&
		a.parallelism == b.parallelism &&
		a.saltLength == b.saltLength &&
		a.keyLength == b.keyLength

}

// Default Argon2id parameters
var defaultArgon2Params = argon2Params{
	memory:      1 << 14, // 64MB
	iterations:  2,       // java 2, golang 3
	parallelism: 4,
	saltLength:  16,
	keyLength:   32,
}

type Argon2Pwd struct {
	p    argon2Params
	salt []byte
	hash []byte
}

func NewArgon2Pwd() PasswordEncoder {
	return &Argon2Pwd{p: defaultArgon2Params, salt: nil}
}

func (a *Argon2Pwd) Encode(rawpass string) (string, error) {
	err := a.onlyHash([]byte(rawpass))
	if err != nil {
		return "", err
	}
	return a.toString(), nil
}

func (a *Argon2Pwd) onlyHash(key []byte) error {
	if a.salt == nil {
		salt := make([]byte, a.p.saltLength)
		r, _ := io.ReadFull(rand.Reader, salt)
		if uint32(r) != a.p.saltLength {
			return SaltFailed
		}
		a.salt = salt[0:]
	}

	a.hash = argon2.IDKey(key, a.salt, a.p.iterations, a.p.memory, a.p.parallelism, a.p.keyLength)
	return nil
}

func (a *Argon2Pwd) Matches(rawpass string, encodedHash string) bool {
	other, err := parseArgon2(encodedHash)
	if err != nil {
		return false
	}

	if !compare(&a.p, &other.p) {
		return false
	}

	// 用相同的salt
	a.salt = make([]byte, a.p.saltLength)
	copy(a.salt, other.salt)
	_ = a.onlyHash([]byte(rawpass))

	return subtle.ConstantTimeCompare(a.hash, other.hash) == 1
}

func (a *Argon2Pwd) toString() string {
	// Encoding to base64 format
	b64Salt := base64.RawStdEncoding.EncodeToString(a.salt)
	b64Hash := base64.RawStdEncoding.EncodeToString(a.hash)

	// Format: $argon2id $v = 19 $m = 65536, t = 3, p = 4 $< salt > $< hash >
	encoded := fmt.Sprintf(
		"$argon2id$v=19$m=%d,t=%d,p=%d$%s$%s",
		a.p.memory, a.p.iterations, a.p.parallelism, b64Salt, b64Hash)

	return encoded
}

func parseArgon2(encodedHash string) (*Argon2Pwd, error) {
	// Parse the encoded hash string
	parts := strings.Split(encodedHash, "$")
	if len(parts) != 6 {
		return nil, ParseFailed
	}

	var p argon2Params
	_, err := fmt.Sscanf(parts[3], "m=%d,t=%d,p=%d", &p.memory, &p.iterations, &p.parallelism)
	if err != nil {
		return nil, ParseFailed
	}

	salt, err := base64.RawStdEncoding.DecodeString(parts[4])
	if err != nil {
		return nil, ParseFailed
	}
	p.saltLength = uint32(len(salt))

	decodedHash, err := base64.RawStdEncoding.DecodeString(parts[5])
	if err != nil {
		return nil, ParseFailed
	}
	p.keyLength = uint32(len(decodedHash))

	return &Argon2Pwd{p: p, salt: salt, hash: decodedHash}, nil
}
